import argparse
import os
import numpy as np
# 导入自定义工具
from utils.h5data import read_h5_data
from utils.phy import *

def plot_dvv_vs_common( te_bjt, dvv, te_phy,phy,fit_quality,
                       output_file, phy_name,title):
    
    ne = len(te_bjt)
    # 拟合过程
    t_shift = np.arange(-0.5,0.5,0.01)
    C = np.zeros(len(t_shift))
    R2= np.zeros(len(t_shift))
    phy_shifted = []
    for i,shift_i in  enumerate(t_shift):
        phy_i = merge_data(te_phy,phy,t_merged=te_bjt+shift_i)
        C[i] = np.sum((phy_i-phy_i.mean())*dvv)/ne

        coffs = np.polyfit(phy_i,dvv,1)
        dvv_rms = dvv- np.polyval(coffs,phy_i)
        ss_res = np.sum(dvv_rms ** 2)
        ss_tot = np.sum((dvv - np.mean(dvv)) ** 2)
        R2_i = 1 - (ss_res / ss_tot) if ss_tot != 0 else 0

        phy_shifted.append(phy_i)
        R2[i] = R2_i
    C = np.abs(C)
    best_idx = np.argmax(R2)
    # best_idx = np.argmax(C)
    shift_best = t_shift[best_idx]
    phy_best = phy_shifted[best_idx]
    R2_best = R2[best_idx]
    shift_best = t_shift[best_idx]

    
    coffs = np.polyfit(phy_best,dvv,1)
    dvv_fit = np.polyval(coffs,phy_best)
    dvv_rms = dvv-dvv_fit
    

    from pylab import figure, rcParams, plt
    rcParams['font.family'] = 'Arial'
    rcParams['font.size'] = '8'
    

    # 图1 绘制搜索过程
    fig = figure(figsize=(3,2.5), dpi=500)
    plt.plot(t_shift*24,R2, '-', color='blue',linewidth=0.6)
    plt.xlabel('time shift(hours)', color='blue')
    plt.ylabel('R', color='blue')
    output_fileNew = output_file.replace('.png', f'.{phy_name}.search.png')
    plt.savefig(output_fileNew, dpi=600, bbox_inches='tight')
    plt.close('all')

    # 图2 绘制最优拟合曲线的线性关系
    fig = figure(figsize=(3,3), dpi=500)
    plt.scatter(phy_best,dvv, marker='^', color='gray',
                linewidth=0.6, s=3, edgecolors='none')
    plt.scatter(phy_best,dvv_fit, marker='o', color='red',
                linewidth=0.6, s=3, edgecolors='none')
    plt.xlabel(f'{phy_name}(dt={shift_best*24:.1f}h)', color='blue')
    plt.ylabel(f'dvv(%)', color='blue')
    plt.title(f'R²={R2_best:.2f}')
    output_fileNew = output_file.replace('.png', f'.{phy_name}.best.png')
    plt.savefig(output_fileNew, dpi=600, bbox_inches='tight')
    plt.close('all')

    # 图3 绘制拟合效果
    fig = figure(figsize=(8, 2.5), dpi=500)
    ax1 = plt.subplot(1, 1, 1)
    print(te_bjt[0],te_bjt[-1])
    plt.errorbar(te_bjt, dvv_fit, yerr=dvv_err*3, 
                 fmt='.-', linewidth=0.7,color='k', markersize=3, markeredgecolor='none',
                ecolor='none',elinewidth=0.6, capsize=1, capthick=0.6, barsabove=True
                     )
    plt.errorbar(te_bjt, dvv, yerr=dvv_err*3, 
                 fmt='.-', linewidth=0.7,color='tab:red', markersize=3, markeredgecolor='none',
                ecolor='none',elinewidth=0.6, capsize=1, capthick=0.6, barsabove=True
                     )
    plt.fill_between(te_bjt, dvv-dvv_err*3, dvv+dvv_err*3, alpha=0.4, color='tab:red')
    plt.xlabel('day (BJT)')
    plt.ylabel('dv/v(%)',color='tab:red')
    plt.ylim([dvv.min()-np.std(dvv)*1,dvv.max()+np.std(dvv)])

    ax2 = plt.twinx(ax1)
    plt.plot(te_bjt, phy_best, '--', color='green',linewidth=0.6)
    plt.ylabel(f'{phy_name}(dt={shift_best*24:.1f}h)', color='green')
    plt.ylim([phy_best.min()-np.std(phy_best), phy_best.max()+np.std(phy_best)])

    plt.title(title)
    plt.grid(True, alpha=0.3)
    
    output_fileNew = output_file.replace('.png', f'.{phy_name}.fit.png')
    plt.savefig(output_fileNew, dpi=600, bbox_inches='tight')
    plt.close('all')

    return dvv_fit, phy_best,shift_best

def plot_dvv_VS_st(phy_file, te_bjt, dvv, dvv_err=None, fit_quality=None, output_file=None,title=''):
    '''
    绘制dv/v随时间变化的曲线
    '''
    # 保存的数据格式：data_in_dict={'te_bjt':te_bjt,
                            #    'slopes':slopes,
                            #    'R2':R2,
                            #    'phy_all':phy_all,
                            #    'phy_mean':phy_mean,
                            #    'STATS':np.array(STATS, dtype='S'),
                            #    'START_TIME':START_TIME},
    # 温度拟合
    # file = 'loc/2406.2406.sensorT.h5'
    # 
    
    te_phy = read_h5_data(phy_file,['te_bjt'])[0]
    phy = read_h5_data(phy_file,['phy_mean'])[0]
    START_TIMEF = read_h5_data(phy_file,['START_TIME'])[0].decode()
    name = 'st2406'

    print(START_TIMEF)

    dvv_fit, phy_best,shift_best = plot_dvv_vs_common(te_bjt, dvv.copy(), 
                                   te_phy,phy,fit_quality,
                                   output_file, name,title=title)

    return dvv-dvv_fit 

def plot_dvv_VS_tide(te_bjt, dvv, dvv_err=None, fit_quality=None, output_file=None,title=''):
    '''
    绘制dv/v随时间变化的曲线
    '''

    # 温度拟合
    nd = te_bjt.max()+1
    te_phy, phy = get_solid(START_TIME, nd)

    dvv_fit, phy_best,shift_best = plot_dvv_vs_common(te_bjt, dvv, 
                                   te_phy,phy,fit_quality,
                                   output_file, 'tide',title=title)

    return dvv-dvv_fit

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Plot traceSeq h5 data')
    parser.add_argument('-input', default='', help='input h5 file')
    parser.add_argument('-figroot', default='figures/7.dt.change.figures', help='root to save figs')
    # parser.add_argument('-figroot', default='figures/debug', help='root to save figs')
    args = parser.parse_args()
    INPUT_FILE = args.input
    FIG_ROOT = args.figroot

    # save_h5_data(OUTFILE, 
    #             data_in_dict={'te_bjt':te_bjt, 'dvv':dvv, 'dvv_err':dvv_err, 'dt_err':dt_err, 'fit_quality':fit_quality,
    #                         'all_groups':all_groups,'MARKER':MARKER, 'EMARKER':EMARKER, 'DATE':DATE},)
    # 读出全部数据
    te_bjt = read_h5_data(INPUT_FILE, ['te_bjt'])[0]
    dvv = read_h5_data(INPUT_FILE, ['dvv'])[0]
    dvv_err = read_h5_data(INPUT_FILE, ['dvv_err'])[0]
    dt_err = read_h5_data(INPUT_FILE, ['dt_err'])[0]
    fit_quality = read_h5_data(INPUT_FILE, ['fit_quality'])[0]
    MARKER = read_h5_data(INPUT_FILE, ['MARKER'])[0].decode()
    DATE = read_h5_data(INPUT_FILE, ['DATE'])[0].decode()
    EMARKER = read_h5_data(INPUT_FILE, ['EMARKER'])[0].decode()

    if DATE=='2406':
        phy_file = 'loc/2406.P264S.sensorT.h5'
        START_TIME = '2024-05-30T00:00:00Z'

    if DATE=='254C':
        phy_file = 'loc/254C.254C.sensorT.h5'
        START_TIME = '2024-10-08T00:00:00Z'


    FIG_ROOT = f'{FIG_ROOT}/7C.{DATE}.{EMARKER}'
    if not os.path.exists(FIG_ROOT):
        os.makedirs(FIG_ROOT)

    mask = np.where(np.all([te_bjt>=65,te_bjt<=75],axis=0))
    # mask = np.where(np.all([te_bjt<=100]))
    dvv     = -dvv[mask]/10 # 0.001转换为百分比,并校正X的符号
    dvv_err = dvv_err[mask]/10
    dt_err  = dt_err[mask]
    te_bjt  = te_bjt[mask]
    fit_quality= fit_quality[mask]
    print(te_bjt)

    # 绘制dv/v曲线
    dvv=plot_dvv_VS_st(phy_file,te_bjt, dvv, dvv_err, fit_quality,
                            title=f'{MARKER} DV/V vs Time',
                        output_file=f'{FIG_ROOT}/dvv_vs_phy.{MARKER}.png')
    # print(dvv)
    # plot_dvv_VS_tide(te_bjt, dvv, dvv_err, fit_quality,
    #                         title=f'{MARKER} DV/V vs Time',
    #                     output_file=f'{FIG_ROOT}/dvv_vs_phy.{MARKER}.png')